
from app.db.models.asset_tasks import AssetTasks
from app.db.models.tasks import Task
from app.db.models.wass import ScanEvents
from app.libs.score import Score
from app.libs.oms import Client
from app.libs.enums import WS_TASK_GROUP_TUPLE, TaskTriggerType, AssetType


def get_host_vul_task_list(query, page, count, sort, status=None, query_level=None, is_web=None):
    # 1、查询任务列表
    data = AssetTasks.objects.find(query)
    total = data.count()

    if sort:
        data = data.order_by(*sort)

    if (page - 1) * count >= total:
        return total, []

    data = data.paginate(page, count)
    asset_task_list = list(data.items)

    # 2、查询任务进度
    schedule_task_id_list = []
    manual_task_id_list = []
    for asset_task in asset_task_list:
        if asset_task.triggerType == TaskTriggerType.schedule.value:
            if asset_task.hostVulSettings.taskId:
                schedule_task_id_list.append(asset_task.hostVulSettings.taskId)
        else:
            if asset_task.hostVulSettings.taskId:
                manual_task_id_list.append(asset_task.hostVulSettings.taskId)

    manual_tasks = Task.objects.filter(taskId__in=manual_task_id_list)
    manual_task_maps = {task.taskId: task.to_mongo().to_dict() for task in manual_tasks}

    pipeline = ([
        {'$match': {'taskId': {"$in": schedule_task_id_list}}},
        {'$group': {'_id': '$taskId', 'id': {"$last": "$_id"}, 'startTime': {"$last": "$startTime"},
                    'nextExecDate': {"$last": "$nextExecDate"}, 'endTime': {"$last": "$endTime"},
                    'status': {"$last": "$status"},
                    'progress': {'$last': '$progress'}}},
    ])
    schedule_tasks = Task.objects.aggregate(pipeline)
    schedule_task_maps = {task.get("_id"): task for task in schedule_tasks}

    results = []
    for asset_task in asset_task_list:
        vul = {'low': 0, 'middle': 0, 'high': 0, 'score': 0, 'level': '未知'}
        last_result = asset_task.lastResult
        if last_result:
            vul.update(last_result.get('host_vul', {}))
            if 'medium' in vul:
                vul['middle'] = vul.get('medium', 0)
                del vul['medium']

        task_id = asset_task.hostVulSettings.taskId

        if asset_task.triggerType == TaskTriggerType.schedule.value:
            task = schedule_task_maps.get(task_id)
        else:
            task = manual_task_maps.get(task_id)
            if task:
                task['id'] = str(task['_id'])

        item = {
            'uid': str(asset_task.uid),
            'task_id': task['id'] if task else '',
            'result_id': str(asset_task.hostVulResultId) if asset_task.hostVulResultId else '',
            'asset_task_id': str(asset_task.pk),
            'job_id': str(asset_task.jobId),
            'name': asset_task.name,
            'target_url': asset_task.targetUrl,
            'trigger_type': asset_task.triggerType,
            'enable_monitor': asset_task.enableMonitor,
            'start_time': task.get('startTime') if task and task.get('startTime') else '',
            'next_time': task.get('nextExecDate') if task and task.get('nextExecDate') else '',
            'update_time': task.get('endTime') if task and task.get('endTime') else '',
            'detail': {
                'process': {"progress": task.get('progress', 0), "status": task.get('status', 'unkonwn')} if task else {"progress": 0, "status": 'unkonwn'},
                'vul': vul,
            },
        }
        results.append(item)

    if status or query_level:
        total = len(results)
        results = results[(page - 1) * count: page * count]

    return total, results


def update_last_web_security_status(job):
    update_dict = {'score': 0, "level": "未知", "update_time": ""}
    if job.assetType == AssetType.host.value:
        return
    all_warnings = []
    end_times = []
    target_status_list = []
    for typ in WS_TASK_GROUP_TUPLE:
        if task := Task.objects.filter(jobId=job.id, taskType=typ, status="completed").order_by("-endTime").first():
            if target_status := task.result.targetStatus:
                target_status_list.append(target_status.status)
            else:
                target_status_list.append("good")
            all_warnings.extend(task.result.warnings)
            end_times.append(task.endTime)
            setattr(job, f"{typ}ResultId", task.id)
        else:
            setattr(job, f"{typ}ResultId", None)
    update_time = max(end_times) if end_times else ''
    if all_warnings:
        score_dict = Score.score_info(event_map={}, warnings=all_warnings)
        update_dict['score'] = score_dict.get('score', 0)
        update_dict['level'] = score_dict.get('level', "未知")

    if (not all_warnings and update_time) and (not target_status_list or "good" in target_status_list):
        update_dict['level'] = '安全'
    elif (not all_warnings and update_time) and (target_status_list and "good" not in target_status_list):
        update_dict['level'] = "未知"

    update_dict['update_time'] = update_time
    job.securityStatus = update_dict
    job.save()


def update_last_host_security_status(job, task=None):
    update_dict = {'score': 0, "level": "未知", "update_time": ""}
    if not task:
        task = Task.objects.filter(jobId=job.id, taskType="hostVul", status="completed").order_by("-id").first()
    if task:
        query = {"task_session_id": task.taskSessionId, "event_name": "network_vuln"}
        host_vul_events = ScanEvents.objects.find(query)
        all_warnings = []
        for event in host_vul_events:
            event.__setattr__('level', event.severity)
            all_warnings.append(event)
        if all_warnings:
            score_dict = Score.score_info(event_map={}, warnings=all_warnings)
            update_dict['score'] = score_dict.get('score', 0)
            update_dict['level'] = score_dict.get('level', "未知")

        update_time = task.endTime or ''
        if not all_warnings and update_time != '':
            update_dict['level'] = '安全'
        update_dict['update_time'] = task.endTime
    job.securityStatus = update_dict
    job.save()


def get_cloud_assets(user):
    oms_client = Client()
    cloud_assets = oms_client.get('/api/inner/', params={
        'type': 'user_property_list',
        'username': user.username
    }).get('ret', [])
    return cloud_assets
